package com.temp.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.*;

/**
 * 优化的WebSocket服务端点 - 修复心跳调度器版本
 */
@Component
@ServerEndpoint("/chat/{userId}")
@Slf4j
public class WebSocketEndpoint {

    private static final ConcurrentHashMap<Long, SessionInfo> userSessions = new ConcurrentHashMap<>();

    // 使用更可靠的调度器创建方式
    private static final ScheduledExecutorService scheduler = createHeartbeatScheduler();

    private static final long HEARTBEAT_INTERVAL = 30;
    private static final long HEARTBEAT_TIMEOUT = 60;
    private static final String PING_MESSAGE = "PING";

    /**
     * 创建可靠的心跳调度器
     */
    private static ScheduledExecutorService createHeartbeatScheduler() {
        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(
                2, // 核心线程数
                new ThreadFactory() {
                    private int threadCount = 0;

                    @Override
                    public Thread newThread(Runnable r) {
                        Thread thread = new Thread(r, "ws-heartbeat-" + (threadCount++));
                        thread.setDaemon(true); // 设置为守护线程
                        thread.setPriority(Thread.NORM_PRIORITY);
                        return thread;
                    }
                },
                new ThreadPoolExecutor.CallerRunsPolicy() // 拒绝策略
        );

        // 配置调度器参数
        executor.setRemoveOnCancelPolicy(true); // 取消后立即移除任务
        executor.setContinueExistingPeriodicTasksAfterShutdownPolicy(false); // 关闭后不继续周期任务
        executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false); // 关闭后不执行延迟任务

        log.info("✅ 心跳调度器创建成功: 核心线程数={}, 线程工厂={}",
                executor.getCorePoolSize(), executor.getThreadFactory());
        return executor;
    }

    // 静态初始化块测试调度器
    static {
        log.info("WebSocketEndpoint类加载完成");

        // 立即测试调度器是否工作
        scheduler.schedule(() -> {
            log.info("✅ 调度器初始化测试：心跳调度器正常工作");
        }, 1, TimeUnit.SECONDS);

        // 添加关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            log.info("正在关闭心跳调度器...");
            scheduler.shutdown();
            try {
                if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                    scheduler.shutdownNow();
                }
            } catch (InterruptedException e) {
                scheduler.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }));
    }

    /**
     * 会话信息封装类
     */
    private static class SessionInfo {
        Session session;
        ScheduledFuture<?> heartbeatFuture;
        volatile long lastActiveTime;
        Long userId;

        SessionInfo(Session session, ScheduledFuture<?> heartbeatFuture, Long userId) {
            this.session = session;
            this.heartbeatFuture = heartbeatFuture;
            this.userId = userId;
            this.lastActiveTime = System.currentTimeMillis();
        }

        void updateActiveTime() {
            this.lastActiveTime = System.currentTimeMillis();
        }

        void cancelHeartbeat() {
            if (heartbeatFuture != null && !heartbeatFuture.isCancelled()) {
                heartbeatFuture.cancel(true);
                log.debug("用户[{}]心跳任务已取消", userId);
            }
        }
    }

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") Long userId) {
        log.info("用户[{}]开始连接", userId);

        // 设置会话超时
        session.setMaxIdleTimeout(HEARTBEAT_INTERVAL * 1000 * 2);

        // 关闭旧连接（原子操作）
        SessionInfo oldInfo = userSessions.remove(userId);
        if (oldInfo != null) {
            log.info("关闭用户[{}]的旧连接", userId);
            oldInfo.cancelHeartbeat();
            cleanupSession(oldInfo.session, userId);
        }

        // 创建新定时任务
        final Long finalUserId = userId;
        ScheduledFuture<?> heartbeatFuture = scheduler.scheduleAtFixedRate(() -> {
            heartbeatTask(finalUserId);
        }, 5, HEARTBEAT_INTERVAL, TimeUnit.SECONDS); // 首次延迟5秒

        // 创建并保存SessionInfo
        SessionInfo newInfo = new SessionInfo(session, heartbeatFuture, userId);
        userSessions.put(userId, newInfo);

        log.info("用户[{}]连接完成，心跳任务已启动，当前在线用户数: {}", userId, userSessions.size());

        // 立即记录调度器状态
        log.debug("调度器活跃线程数: {}", ((ThreadPoolExecutor) scheduler).getActiveCount());
    }

    /**
     * 独立的心跳任务方法
     */
    private void heartbeatTask(Long userId) {
        try {
            log.info("❤️ 心跳任务执行 - 用户[{}]", userId);

            SessionInfo info = userSessions.get(userId);
            if (info == null) {
                log.warn("用户[{}]会话信息为null，停止心跳", userId);
                return;
            }
            if (!info.session.isOpen()) {
                log.warn("用户[{}]会话已关闭，停止心跳", userId);
                return;
            }

            long inactiveTime = System.currentTimeMillis() - info.lastActiveTime;
            log.debug("用户[{}]空闲时间: {}ms, 超时阈值: {}ms",
                    userId, inactiveTime, HEARTBEAT_TIMEOUT * 1000);

            // 检查超时
            if (inactiveTime > HEARTBEAT_TIMEOUT * 1000) {
                log.warn("用户[{}]心跳超时，断开连接", userId);
                throw new IOException("心跳超时：" + inactiveTime + "ms");
            }

            // 发送Ping
            info.session.getBasicRemote().sendPing(ByteBuffer.wrap(PING_MESSAGE.getBytes()));
            log.info("✅ 向用户[{}]发送心跳PING", userId);

        } catch (Exception e) {
            log.error("❌ 用户[{}]心跳检测失败: {}", userId, e.getMessage());
            SessionInfo errorInfo = userSessions.get(userId);
            if (errorInfo != null) {
                cleanupSession(errorInfo.session, userId);
            }
        }
    }

    @OnClose
    public void onClose(Session session, @PathParam("userId") Long userId) {
        log.info("用户[{}]连接关闭", userId);
        SessionInfo info = userSessions.get(userId);
        if (info != null && info.session.getId().equals(session.getId())) {
            info.cancelHeartbeat();
            cleanupSession(session, userId);
            log.info("用户[{}]已断开, 当前在线用户数: {}", userId, userSessions.size());
        }
    }

    @OnError
    public void onError(Session session, @PathParam("userId") Long userId, Throwable error) {
        log.error("用户[{}]连接错误: {}", userId, error.getMessage(), error);
        SessionInfo info = userSessions.get(userId);
        if (info != null && info.session.getId().equals(session.getId())) {
            info.cancelHeartbeat();
            cleanupSession(session, userId);
        }
    }

    @OnMessage
    public void onMessage(String message, Session session, @PathParam("userId") Long userId) {
        log.debug("收到用户[{}]的消息: {}", userId, message);
        SessionInfo info = userSessions.get(userId);
        if (info != null) {
            info.updateActiveTime();
            if ("PING".equals(message)) {
                log.debug("收到用户[{}]的心跳PING请求", userId);
                try {
                    session.getBasicRemote().sendText("PONG");
                    log.debug("向用户[{}]发送PONG响应", userId);
                } catch (IOException e) {
                    log.error("发送PONG响应失败", e);
                }
            } else {
                log.debug("用户[{}]发送普通消息: {}", userId, message);
                if (AdminWebSocketEndpoint.isAdminOnline()) {
                    AdminWebSocketEndpoint.sendToAdmin(userId, message);
                }
            }
        }
    }

    public static Session getUserSession(Long userId) {
        SessionInfo info = userSessions.get(userId);
        return info != null ? info.session : null;
    }

    public static boolean isUserOnline(Long userId) {
        SessionInfo info = userSessions.get(userId);
        return info != null && info.session.isOpen();
    }

    private static void cleanupSession(Session session, Long userId) {
        try {
            if (session != null && session.isOpen()) {
                session.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "连接关闭"));
            }
        } catch (IOException e) {
            log.error("关闭用户[{}]会话失败: {}", userId, e.getMessage());
        } finally {
            SessionInfo currentInfo = userSessions.get(userId);
            if (currentInfo != null && currentInfo.session.getId().equals(session.getId())) {
                userSessions.remove(userId);
                log.info("用户[{}]会话清理完成", userId);
            }
        }
    }

    /**
     * 获取调度器状态信息（用于监控）
     */
    public static String getSchedulerStatus() {
        ThreadPoolExecutor executor = (ThreadPoolExecutor) scheduler;
        return String.format(
                "调度器状态: 活跃线程=%d, 池大小=%d, 队列大小=%d, 已完成任务=%d",
                executor.getActiveCount(),
                executor.getPoolSize(),
                executor.getQueue().size(),
                executor.getCompletedTaskCount()
        );
    }
}